package main

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"os"
)

// 生成一个随机的 32 字节密钥（AES-256）
func generateKey() []byte {
	key := make([]byte, 32)
	if _, err := rand.Read(key); err != nil {
		panic(err)
	}
	return key
}

// 生成一个随机的 16 字节 IV（AES 的 BlockSize 是 16）
func generateIV() []byte {
	iv := make([]byte, aes.BlockSize)
	if _, err := rand.Read(iv); err != nil {
		panic(err)
	}
	return iv
}

// PKCS7 填充，保证数据是 16 字节的倍数
func pkcs7Pad(data []byte, blockSize int) []byte {
	padding := blockSize - (len(data) % blockSize)
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(data, padText...)
}

// PKCS7 去除填充
func pkcs7Unpad(data []byte) ([]byte, error) {
	var length, padding int
	if length = len(data); length == 0 {
		return nil, errors.New("invalid padding size")
	}
	padding = int(data[length-1])
	if padding > length {
		return nil, errors.New("invalid padding content")
	}
	return data[:length-padding], nil
}

// AES-256 CBC 加密
func encryptAES256CBC(plaintext, key, iv []byte) ([]byte, error) {
	block, err := aes.NewCipher(key) // 创建 AES 密码块
	if err != nil {
		return nil, err
	}

	plaintext = pkcs7Pad(plaintext, aes.BlockSize)

	ciphertext := make([]byte, len(plaintext))
	// 这个实例封装了 AES 算法 + CBC 模式 + IV（加盐），用于加密数据块
	mode := cipher.NewCBCEncrypter(block, iv)
	mode.CryptBlocks(ciphertext, plaintext)

	return ciphertext, nil
}

// AES-256 CBC 解密
func decryptAES256CBC(ciphertext, key, iv []byte) ([]byte, error) {
	var block cipher.Block
	var err error
	if block, err = aes.NewCipher(key); err != nil {
		return nil, err
	}

	if len(ciphertext)%aes.BlockSize != 0 {
		return nil, errors.New("ciphertext is not a multiple of the block size")
	}

	plaintext := make([]byte, len(ciphertext))
	mode := cipher.NewCBCDecrypter(block, iv)
	mode.CryptBlocks(plaintext, ciphertext)

	return pkcs7Unpad(plaintext)
}

// 将加密数据写入文件
func encryptAndWriteToFile(filename string, plaintext []byte, iv, key []byte) error {

	var file *os.File
	var err error
	var ciphertext []byte
	// 加密数据
	if ciphertext, err = encryptAES256CBC(plaintext, key, iv); err != nil {
		return err
	}

	if file, err = os.Create(filename); err != nil {
		return err
	}
	defer file.Close()

	_, err = file.Write(ciphertext)
	return err
}

// 读取加密文件并解密
func readAndDecryptFile(filename string, iv, key []byte) ([]byte, error) {
	var file *os.File
	var err error
	var ciphertext []byte

	if file, err = os.Open(filename); err != nil {
		return nil, err
	}
	defer file.Close()

	// 读取密文
	if ciphertext, err = io.ReadAll(file); err != nil {
		return nil, err
	}

	// 解密
	return decryptAES256CBC(ciphertext, key, iv)
}

func main() {
	// 生成 AES-256 密钥
	key := generateKey()
	iv := generateIV()
	fmt.Println("使用的密钥（Base64编码）:", base64.StdEncoding.EncodeToString(key))
	fmt.Println("使用的密钥（Base64编码）:", base64.StdEncoding.EncodeToString(iv))

	plaintext := make([]byte, 0)
	// 需要加密的数据
	for i := 0; i < 100; i++ {
		plaintext = append(plaintext, []byte("这是一个需要加密的文件内容！")...)
	}

	// 定义文件名
	filename := "encrypted_data.txt"

	// 加密并写入文件
	if err := encryptAndWriteToFile(filename, plaintext, iv, key); err != nil {
		fmt.Println("加密失败:", err)
	} else {
		fmt.Println("加密并写入文件成功！")
	}

	// 读取并解密
	if decryptedText, err := readAndDecryptFile(filename, iv, key); err != nil {
		fmt.Println("解密失败:", err)
	} else {
		fmt.Println("解密成功，内容:", string(decryptedText))
	}
}
